{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "#随机生成向量测试\n",
    "import numpy as np\n",
    "import time\n",
    "d = 4                                           # 向量维度\n",
    "nb = 10                                      # index向量库的数据量\n",
    "nq = 2                                       # 待检索query的数目\n",
    "np.random.seed(1234)             \n",
    "xb = np.random.random((nb, d)).astype('float32')\n",
    "xb[:, 0] += np.arange(nb) / 1000.                # index向量库的向量\n",
    "xq = np.random.random((nq, d)).astype('float32')\n",
    "xq[:, 0] += np.arange(nq) / 1000.                # 待检索的query向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "base:\n",
      " [[0.19151945 0.62210876 0.43772775 0.7853586 ]\n",
      " [0.7809758  0.2725926  0.27646425 0.8018722 ]\n",
      " [0.96013933 0.87593263 0.35781726 0.5009951 ]\n",
      " [0.68646294 0.71270204 0.37025076 0.5611962 ]\n",
      " [0.5070832  0.01376845 0.7728266  0.8826412 ]\n",
      " [0.36988598 0.6153962  0.07538124 0.368824  ]\n",
      " [0.9391401  0.65137815 0.39720258 0.78873014]\n",
      " [0.32383612 0.56809866 0.8691274  0.4361734 ]\n",
      " [0.81014764 0.14376682 0.70426095 0.7045813 ]\n",
      " [0.22779211 0.92486763 0.44214076 0.90931594]]\n",
      "query:\n",
      " [[0.05980922 0.18428709 0.04735528 0.6748809 ]\n",
      " [0.59562474 0.5333102  0.04332406 0.5614331 ]]\n"
     ]
    }
   ],
   "source": [
    "print(\"base:\\n\", xb[:10])\n",
    "print(\"query:\\n\", xq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "centroids:\n",
      " [[0.19151945 0.62210876 0.43772775 0.7853586 ]\n",
      " [0.7809758  0.2725926  0.27646425 0.8018722 ]\n",
      " [0.96013933 0.87593263 0.35781726 0.5009951 ]\n",
      " [0.68646294 0.71270204 0.37025076 0.5611962 ]\n",
      " [0.5070832  0.01376845 0.7728266  0.8826412 ]\n",
      " [0.36988598 0.6153962  0.07538124 0.368824  ]\n",
      " [0.9391401  0.65137815 0.39720258 0.78873014]\n",
      " [0.32383612 0.56809866 0.8691274  0.4361734 ]\n",
      " [0.81014764 0.14376682 0.70426095 0.7045813 ]\n",
      " [0.22779211 0.92486763 0.44214076 0.90931594]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING clustering 10 points to 10 centroids: please provide at least 390 training points\n"
     ]
    }
   ],
   "source": [
    "import faiss\n",
    "kmeans = faiss.Kmeans(d, 10)\n",
    "kmeans.train(xb)\n",
    "centroids = kmeans.centroids\n",
    "print(\"centroids:\\n\",centroids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([[0.        , 0.10836354, 0.30797812, 0.32846075, 0.33665574],\n",
      "       [0.        , 0.1832449 , 0.20992261, 0.2693498 , 0.39490637],\n",
      "       [0.        , 0.10532176, 0.13520834, 0.4932641 , 0.5135175 ],\n",
      "       [0.        , 0.10532176, 0.12010445, 0.23364449, 0.2693498 ],\n",
      "       [0.        , 0.14515421, 0.39490637, 0.5494688 , 0.59141356],\n",
      "       [0.        , 0.23364449, 0.33665574, 0.5135175 , 0.5144743 ],\n",
      "       [0.        , 0.12010445, 0.13520834, 0.1832449 , 0.3756742 ],\n",
      "       [0.        , 0.32846075, 0.41691694, 0.51578015, 0.54269   ],\n",
      "       [0.        , 0.14515421, 0.20992261, 0.3756742 , 0.47110727],\n",
      "       [0.        , 0.10836354, 0.3817487 , 0.5426073 , 0.54269   ]],\n",
      "      dtype=float32), array([[0, 9, 3, 7, 5],\n",
      "       [1, 6, 8, 3, 4],\n",
      "       [2, 3, 6, 1, 5],\n",
      "       [3, 2, 6, 5, 1],\n",
      "       [4, 8, 1, 7, 0],\n",
      "       [5, 3, 0, 2, 1],\n",
      "       [6, 3, 2, 1, 8],\n",
      "       [7, 0, 3, 8, 9],\n",
      "       [8, 4, 1, 6, 3],\n",
      "       [9, 0, 3, 5, 7]]))\n"
     ]
    }
   ],
   "source": [
    "labels = kmeans.index.search(xb, 5)#[1].reshape(-1)\n",
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
